#####################################################
#-------------------  USAGE  -----------------------#
#####################################################

# CountAAs.pl usage:
#	-o <base output filename> | required, extension COUNTS will be added
#	-i <input masked faa file> | required
#	-a <input original faa file> | required
#	-r <regular expression for line matching> | required - case insensitive, e.g. AAW or X.
#	-c <*.stat file created from cast> | optional, adds a column with the AAs masked for every sequence

	

#####################################################
#-----------------  INCLUDES  ----------------------#
#####################################################

use strict;
use Getopt::Std;
use Bio::Seq;
use Bio::SeqIO;

#####################################################
#---------------  DECLARATIONS  --------------------#
#####################################################

my %options=();
my $count = 0;
my $totalSeq=0;
my $totalSeqFound=0;
my $totalAAs=0;
my $totalAAsFound=0;
my $perAA=0;
my @lines =();
my %stats=();

my $usage = "CountAAs.pl usage:
	-o <base output filename> | required, extension COUNTS will be added
	-i <input file> | required
	-a <input original faa file> | required
	-r <regular expression for line matching> | required - case insensitive, e.g. AAW or X.
	-c <*.stat file created from cast> | optional, adds a column with the AAs masked for every sequence
	";

#####################################################
#---------------  GET OPTIONS  ---------------------#
#####################################################

unless (getopts("o:i:a:r:c:", \%options)) {die "$usage\n";} # wrong arguments provided
unless($options{o} && ($options{i}) && $options{r} && ($options{a})) {die "$usage\n";}

#####################################################
#--------------------- MAIN  -----------------------#
#####################################################


#$/="\n>";
if (defined $options{c}) {
	open(FSTAT, $options{c}) or die "Cannot open stat file\n";
	
	while (<FSTAT>) 
	{
		chomp;
		s/^\>//g;
		/(\S+).+\t(.+)\t(.+)\t(.+)\t(.+)\n/i;
		if (defined $stats{$1}) {
			if (rindex( uc($stats{$1}), uc($2))>=0) {next;}
#			my $tmp = $stats{$1};
			$stats{$1} = join ",", $stats{$1}, $2;
		} else {
			$stats{$1} = $2;
		}
	}
	close(FSTAT);

}
#$/="\\n";

my %seqs=();
#open OFILE, "$options{a}" or die "Cannot open FASTA file $options{a}\n";
 print STDERR "Now reading original fasta file\n" ;

my $oseqin = Bio::SeqIO->new(-format =>"largefasta", -file=> $options{a}, -alphabet =>'protein') or die "Invalid input file\n";
my $ototalSeq = 0;
while((my $seqobj = $oseqin->next_seq())) {
	my $id = $seqobj->display_id;
	my $count = 0;
	my $sequence = $seqobj->seq();
	$count++ while $sequence =~ /$options{r}/ig;
	$seqs{$id} = $count;
	$ototalSeq++;
	print STDERR "Processed $ototalSeq so far...\r";
}
print STDERR "Done reading $ototalSeq sequences from the original fasta file\n" ;
#close OFILE;
print STDERR "Finished reading original file\n";
my $new=1; my $id1='';
#open FFILE, "$options{i}" or die "Cannot open FASTA file $options{i}\n";

my $seqin = Bio::SeqIO->new(-format =>"largefasta", -file=> $options{i}, -alphabet =>'protein') or die "Invalid input file\n";

print STDERR "Now reading masked fasta file\n" ;

      while((my $seqobj = $seqin->next_seq())) {
         $count = 0; $perAA=0;
         my $len = $seqobj->length();
         my $sequence = $seqobj->seq();
		 my $id = $seqobj->display_id;
         $count++ while $sequence =~ /$options{r}/ig;
		 $count -= $seqs{$id};
         if ($count >0) {$totalSeqFound++; $totalAAsFound+=$count;}
         $totalSeq++; $totalAAs+=$len;
      	 $perAA = sprintf("%.2f", ($count/$len)*100);
      	
      	 
      	 
	      	 if (defined $options{c}) {
	         	push (@lines, $id, "\t", $seqobj->length(),"\t", $count, "\t",  $perAA, , "\t", $stats{$id},"\n"); 	 
	      	 } else {
	        	 push (@lines, $id, "\t", $seqobj->length(),"\t", $count, "\t",  $perAA, "\n");
	         }
		print STDERR "Processed $totalSeq so far...\r";
       }
         
print STDERR "Done reading $totalSeq sequences from the masked fasta file\n" ;

open(OUTPUT, ">$options{o}.COUNTS") or die "Can't open output file\n";

        print OUTPUT "##################################################\n";
        print OUTPUT "#- Total statistics for file $options{i}\n";
        print OUTPUT "##################################################\n";	
		print OUTPUT "### Number of sequences in file: \t$totalSeq\n### Number of sequences matching search term: \t$totalSeqFound\n";
        print OUTPUT "### Number of AAs in file: \t$totalAAs\n### Number of AAs matching search term: \t$totalAAsFound\n";
        print OUTPUT "##################################################\n";
        print OUTPUT "#--------------- List of sequences --------------#\n";
        print OUTPUT "##################################################\n";
        
        foreach my $line (@lines) {print OUTPUT $line;}
close(OUTPUT);       
       
       